#!/bin/bash
#TOP_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )/../.." &> /dev/null && pwd )"
CUR_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"


#src                      mirror1                        mirror2
#------------------------------------------------------------------------------------
#cr.l5d.io                l5d.m.daocloud.io              m.daocloud.io/cr.l5d.io
#docker.elastic.co        elastic.m.daocloud.io          m.daocloud.io/docker.elastic.co
#docker.io                docker.m.daocloud.io           m.daocloud.io/docker.io
#gcr.io                   gcr.m.daocloud.io              m.daocloud.io/gcr.io
#ghcr.io                  ghcr.m.daocloud.io             m.daocloud.io/ghcr.io
#k8s.gcr.io               k8s-gcr.m.daocloud.io          m.daocloud.io/k8s.gcr.io
#registry.k8s.io          k8s.m.daocloud.io              m.daocloud.io/registry.k8s.io
#mcr.microsoft.com        mcr.m.daocloud.io              m.daocloud.io/mcr.microsoft.com
#nvcr.io                  nvcr.m.daocloud.io             m.daocloud.io/nvcr.io
#quay.io                  quay.m.daocloud.io             m.daocloud.io/quay.io
#registry.jujucharms.com  jujucharms.m.daocloud.io       m.daocloud.io/registry.jujucharms.com
#rocks.canonical.com      rocks-canonical.m.daocloud.io  m.daocloud.io/rocks.canonical.com

MIRROS=m.daocloud.io
K8S_version=${K8S_version:-1.21.11}
IMG_TXT=$CUR_DIR/images-${K8S_version}.txt
pyellow() {
  echo -e "\033[33m $1 \033[0m"
}

pblue() {
  echo -e "\033[34m $1 \033[0m"
}

get_image_mirror() {
  local image=$1
  local io=${image%%/*}

  case $io in
    k8s.gcr.io|ghcr.io|gcr.io|docker.io|registry.k8s.io)
      local name=${image##*/}
      echo "$MIRROS/$image" ;;
    *) echo "$image"
  esac
}

clean_download() {
    rm -rf $CUR_DIR/images.tgz
	rm -rf $CUR_DIR/images.tar
}

delete_old_version() {
    imageFiles=${IMG_TXT} #,$CUR_DIR/images-chart.txt
	files=$(echo $imageFiles | sed 's/,/ /g')
	for i in $files; do
	  lines=$(cat $i | grep -vE "^#|^[ ]*$")
	  count=$(echo "$lines" | wc -l)
	  cnt=0
	  while read line; do
		if [ -z "$(echo $line | grep -vE '^#|^[ ]*$')" ]; then
		  continue
		fi
		let cnt=cnt+1 
		#set -x
		localImage=${line%@*}
		if [ -z "$remoteRegistry" ]; then
		  remoteImage=${line#*@}
		  remoteTargetImage=$(get_image_mirror $remoteImage)
		else
		  remoteTargetImage=$remoteRegistry/${localImage#*/}
		fi
		#localTargetImage=$localRegistry/${localImage#*/}
		echo $(pyellow $cnt/$count): $(pblue "pull $remoteTargetImage, push $localTargetImage")
		#set -x 
		bash -xc "docker rmi $remoteTargetImage > /dev/null 2>&1"
		#docker tag $remoteTargetImage $localTargetImage
		#docker push $localTargetImage
		#set +x
		echo "    target=$remoteTargetImage [ok]"
		ALL_IMAGES="$ALL_IMAGES $remoteTargetImage"
	  done < $i
	done
}

print_usage() {
    cat << EOF 
Usage: $0 [--force|-f] [--rm|--remove] [--clean|-c] [--downloadonly|-d] 
Options:
  --force|-f: force installation, default force=false
  --rm|--remove: docker rmi images
  --clean|-c: remove images.tgz
  --downloadonly|-d: only docker pull images, do not tar
EOF
}

install_images_online() {
	ALL_IMAGES=""
	imageFiles=${IMG_TXT} #,$CUR_DIR/images-chart.txt
	files=$(echo $imageFiles | sed 's/,/ /g')
	for i in $files; do
	  lines=$(cat $i | grep -vE "^#|^[ ]*$")
	  count=$(echo "$lines" | wc -l)
	  cnt=0
	  while read line; do
		if [ -z "$(echo $line | grep -vE '^#|^[ ]*$')" ]; then
		  continue
		fi
		let cnt=cnt+1 
		set -x
		image=${line}
		if [ -z "$remoteRegistry" ]; then
		  remoteImage=${line#*@}
		  remoteTargetImage=$(get_image_mirror $image)
		else
		  remoteTargetImage=$remoteRegistry/${image#*/}
		fi
		echo $(pyellow $cnt/$count): $(pblue "pull $image")
		#set -x 
		bash -xc "docker pull $remoteTargetImage"
		if [ "$remoteTargetImage" != "$image" ]; then
		    bash -xc "docker tag $remoteTargetImage ${image}"
		fi
		#docker tag $remoteTargetImage $localTargetImage
		#docker push $localTargetImage
		set +x
		echo "    target=$image [ok]"
		ALL_IMAGES="$ALL_IMAGES ${image}"
	  done < $i
	done
}

install_images_offline() {
    
	if [ -f $CUR_DIR/images.tgz ] && ! $DOWNLOAD_ONLY; then
	    set -x
	    tar zxvf $CUR_DIR/images.tgz -C $CUR_DIR/
	    docker load -i $CUR_DIR/images.tar
            set +x
	else
	    install_images_online
	fi
}

#=====================main func ==============================
if [[ $EUID -ne 0 ]]; then
    echo "This script must be run as root"
    exit 1
fi

#清除旧版本的影响，保证command -v命令得到正确的结果
hash -r

ARGS=$#

# Parse command line arguments
while [[ $# -gt 0 ]]
do
    key="$1"
    case $key in
    --online|-o)
        ONLINE=true
        shift
        ;;
    --clean|-c)
        if [[ $ARGS -eq 1 ]]; then
            clean_download
            exit 0
        fi
        AUTOCLEAN=true
        shift
        ;;
    --downloadonly|-d)
        DOWNLOAD_ONLY=true
        shift
        ;;
    --force|-f)
        FORCE=true
        shift
        ;;
    --help)
        print_usage
        exit 1
        ;;
    --rm|--remove)
        echo "docker rmi old images..."
        delete_old_version
        exit 1
        ;;
    -[^-]*)
        for ((i=1; i<${#key}; i++)); do
          case "${key:$i:1}" in
            f) FORCE=true ;;
            c) AUTOCLEAN=true ;;
            d) DOWNLOAD_ONLY=true ;;
            *) print_usage; exit 1 ;;
          esac
        done
        shift    
        ;;
    *)    # unknown option
        print_usage
        exit 1
        ;;
    esac
done

ONLINE=${ONLINE:-false}
FORCE=${FORCE:-false}
DOWNLOAD_ONLY=${DOWNLOAD_ONLY:-false}
AUTOCLEAN=${AUTOCLEAN:-false}

echo "install images: online=$ONLINE force=$FORCE downloadonly=$DOWNLOAD_ONLY autoclean=$AUTOCLEAN"

$FORCE && delete_old_version

if $ONLINE ; then
    install_images_online
else 
    install_images_offline
fi
if $DOWNLOAD_ONLY; then
    set -x
    docker save -o $CUR_DIR/images.tar $ALL_IMAGES
	cd $CUR_DIR/
	tar zcvf $CUR_DIR/images.tgz images.tar
	rm -rf $CUR_DIR/images.tar
	cd -
fi

